import sys
import json
import numpy as np
from PyQt6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                             QHBoxLayout, QPushButton, QSlider, QLabel, QDoubleSpinBox, QStyle)
from PyQt6.QtCore import QTimer, Qt
from PyQt6.QtGui import QColor, QAction, QFont
import pyqtgraph as pg
import pyqtgraph.opengl as gl

from model_manager import ModelManager, ModelType

# --- 自定义 GLView 以支持右键平移 ---
class CustomGLView(gl.GLViewWidget):
    def mouseMoveEvent(self, ev):
        """
        重写鼠标移动事件：
        PyQtGraph 默认: 左键=旋转, 中键=平移
        修改为: 左键=旋转, 右键=平移
        """
        diff = ev.position() - self.mousePos
        self.mousePos = ev.position()
        
        if ev.buttons() == Qt.MouseButton.LeftButton:
            self.orbit(-diff.x(), diff.y())
        elif ev.buttons() == Qt.MouseButton.RightButton:
            # 映射右键到平移 (x, y, z)
            # 注意：cameraPosition() 返回的是相机中心
            # 这里的平移逻辑沿用 pyqtgraph 内部的 pixel 到 world 的映射
            self.pan(diff.x(), diff.y(), 0, relative='view')

        ev.accept()

    def keyPressEvent(self, ev):
        """
        捕获键盘事件并传递给父窗口(VisualizerWindow)处理
        这样即使焦点在3D视图上，空格键也能控制播放
        """
        if self.parent():
            # 将事件传递给父窗口的 keyPressEvent
            self.parent().keyPressEvent(ev)
        else:
            super().keyPressEvent(ev)
            
# --- 3D 纸飞机 Mesh 生成 (带颜色参数) ---
def create_paper_plane_mesh(size=20, color=(0, 1, 1, 1)):
    """
    创建一个简单的纸飞机形状
    color: (r, g, b, a) 0.0-1.0
    """
    # 顶点坐标 (x, y, z) - 机头朝向 +X
    verts = np.array([
        [2, 0, 0],    # 0: 机头
        [-1, 1, 0],   # 1: 左翼
        [-1, -1, 0],  # 2: 右翼
        [-1, 0, 0.5], # 3: 脊背
        [-1, 0, -0.2] # 4: 腹部
    ], dtype=np.float32) * size

    faces = np.array([
        [0, 1, 3], [0, 3, 2], # 上表面
        [0, 2, 4], [0, 4, 1], # 下表面
        [1, 2, 3], [1, 4, 2]  # 尾部
    ])
    
    # 为每个面应用传入的颜色，稍微做一点明暗变化增强立体感
    base_color = np.array(color)
    face_colors = []
    for i in range(len(faces)):
        # 简单的随机变暗一点点，区分不同面
        shade_factor = 1.0 - (i % 3) * 0.1 
        face_colors.append(base_color * shade_factor)
    
    face_colors = np.array(face_colors, dtype=np.float32)
    # 确保 alpha 通道正确
    face_colors[:, 3] = 1.0

    mesh = gl.GLMeshItem(vertexes=verts, faces=faces, faceColors=face_colors, 
                         smooth=False, shader='shaded', glOptions='opaque')
    return mesh

def create_node(model_type, face_color, edge_color, scale=20.0):
    md = ModelManager.get_mesh_data(model_type)
    mesh = gl.GLMeshItem(
                meshdata=md, 
                smooth=False, 
                #drawEdges=True, 
                #edgeColor=edge_color, 
                shader='shaded', 
                color=face_color)
    mesh.scale(scale, scale, scale)
    return mesh

class VisualizerWindow(QMainWindow):
    def __init__(self, log_file):
        super().__init__()
        self.setWindowTitle("Netlib Visualizer")
        self.resize(1200, 800)
        self.setAcceptDrops(True) # 允许拖拽文件(预留)

        # 3. 播放控制
        self.timer = QTimer()
        self.timer.timeout.connect(self.update_frame)
        self.current_time = 0.0
        self.is_playing = False
        self.playback_speed = 1.0
        self.fps = 30
        
        # 颜色盘 (高亮显示)
        self.node_colors = [
            (0.0, 1.0, 1.0, 1.0), # Cyan
            (1.0, 0.0, 1.0, 1.0), # Magenta
            (0.5, 1.0, 0.0, 1.0), # Lime
            (1.0, 0.5, 0.0, 1.0), # Orange
        ]
        
        # 1. 加载数据
        self.load_data(log_file)

        # 2. 界面初始化
        self.init_ui()
        self.init_scene()



    def load_data(self, filename):
        print("正在加载数据...")
        try:
            with open(filename, 'r') as f:
                data = json.load(f)
        except FileNotFoundError:
            print("错误: 找不到日志文件")
            sys.exit(1)
        
        self.pos_data = {} 
        self.comm_data = [] 
        times = set()

        for entry in data:
            if entry['type'] == 'pos':
                nid = entry['node']
                if nid not in self.pos_data:
                    self.pos_data[nid] = {'times': [], 'pos': []}
                self.pos_data[nid]['times'].append(entry['time'])
                self.pos_data[nid]['pos'].append(entry['pos'])
                times.add(entry['time'])
            elif entry['type'] == 'packet':
                self.comm_data.append(entry)
        
        self.nodes_ids = sorted(self.pos_data.keys())
        for nid in self.nodes_ids:
            self.pos_data[nid]['times'] = np.array(self.pos_data[nid]['times'])
            self.pos_data[nid]['pos'] = np.array(self.pos_data[nid]['pos'])

        self.max_time = max(times) if times else 0
        print(f"数据加载完成. 总时长: {self.max_time:.1f}s, 节点数: {len(self.nodes_ids)}")

    def init_ui(self):
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)

        # --- 3D 视图 ---
        # 使用自定义 View 以支持右键平移
        self.view = CustomGLView()
        self.view.opts['distance'] = 1500
        self.view.setBackgroundColor((100, 100, 100)) # 深灰色背景
        
        # 地面网格
        g = gl.GLGridItem()
        g.scale(1000, 1000, 1)
        self.view.addItem(g)
        
        # 坐标轴
        axis = gl.GLAxisItem()
        axis.setSize(1000, 1000, 1000)
        self.view.addItem(axis)
        
        layout.addWidget(self.view, 1)

        # --- 播放控制栏 ---
        controls = QHBoxLayout()
        
        # 后退按钮
        self.btn_rw = QPushButton("<< 5s")
        self.btn_rw.clicked.connect(lambda: self.seek_relative(-5))
        controls.addWidget(self.btn_rw)

        # 播放/暂停
        self.btn_play = QPushButton()
        # 设置图标
        self.update_play_button_icon()
        self.btn_play.clicked.connect(self.toggle_play)
        controls.addWidget(self.btn_play)

        # 快进按钮
        self.btn_ff = QPushButton("5s >>")
        self.btn_ff.clicked.connect(lambda: self.seek_relative(5))
        controls.addWidget(self.btn_ff)

        # 进度条
        self.slider = QSlider(Qt.Orientation.Horizontal)
        self.slider.setRange(0, int(self.max_time * 1000))
        self.slider.sliderPressed.connect(self.slider_pressed)
        self.slider.sliderReleased.connect(self.slider_released)
        self.slider.valueChanged.connect(self.slider_moved)
        controls.addWidget(QLabel("进度:"))
        controls.addWidget(self.slider)

        self.lbl_time = QLabel("0.000 s")
        self.lbl_time.setFixedWidth(60)
        controls.addWidget(self.lbl_time)

        # 倍速
        controls.addWidget(QLabel("倍速:"))
        self.spin_speed = QDoubleSpinBox()
        self.spin_speed.setDecimals(3)
        self.spin_speed.setRange(0.001, 10.0)
        self.spin_speed.setValue(1.0)
        self.spin_speed.setSingleStep(0.1)
        self.spin_speed.valueChanged.connect(self.set_speed)
        controls.addWidget(self.spin_speed)

        layout.addLayout(controls)

        # 帮助提示
        self.statusBar().showMessage("操作: [空格]=播放/暂停, [左/右键]=逐帧进退，[右键拖动]=平移, [左键拖动]=旋转, [滚轮]=缩放")

    def init_scene(self):
        self.text_font = QFont()
        self.text_font.setPointSize(12)
        self.text_font.setBold(True)

        self.node_items = {}
        self.node_labels = {}  # 存储节点下方的描述文字

        # 初始化节点模型
        for i, nid in enumerate(self.nodes_ids):
            # 获取颜色
            color = self.node_colors[i % len(self.node_colors)]
            # 创建节点
            #mesh = create_paper_plane_mesh(size=30, color=color) # 放大一点尺寸
            #face_color = (0.1, 0.4, 0.8, 0.3)
            face_color = (0.4, 0.9, 1.0, 1.0)
            edge_color = (0.4, 0.9, 1.0, 1.0)
            mesh = create_node(ModelType.UAV_PROPELLER, face_color, edge_color)
            self.view.addItem(mesh)
            self.node_items[nid] = mesh
            # 创建描述标签
            label = gl.GLTextItem(pos=np.array([0,0,0]), text=f'{nid}', 
                color=(255,255,255,255), font=self.text_font)
            self.view.addItem(label)
            self.node_labels[nid] = label

        self.active_packets = {}

    def keyPressEvent(self, event):
        """键盘事件处理"""
        key = event.key()
        if key == Qt.Key.Key_Space:
            self.toggle_play()
        elif key == Qt.Key.Key_Left:
            self.step_frame(-1)
        elif key == Qt.Key.Key_Right:
            self.step_frame(1)
        else:
            super().keyPressEvent(event)

    def step_frame(self, direction):
        """
        逐帧步进
        direction: 1(前进) -1(后退)
        """
        # 如果正在播放，先暂停，以方便用户逐帧观察
        if self.is_playing:
            self.toggle_play()

        # 计算单帧时间步长
        frame_step = (1.0 / self.fps) * self.playback_speed
        
        new_time = self.current_time + (frame_step * direction)
        
        # 边界检查
        new_time = max(0.0, min(new_time, self.max_time))
        
        self.current_time = new_time
        self.update_view()
        self.update_ui_slider()

    def update_play_button_icon(self):
        icon = self.style().standardIcon(
            QStyle.StandardPixmap.SP_MediaPause if self.is_playing 
            else QStyle.StandardPixmap.SP_MediaPlay
        )
        self.btn_play.setIcon(icon)
        self.btn_play.setText(" 暂停 " if self.is_playing else " 播放 ")

    def toggle_play(self):
        self.is_playing = not self.is_playing
        if self.is_playing:
            self.timer.start(int(1000/self.fps))
        else:
            self.timer.stop()
        self.update_play_button_icon()

    def seek_relative(self, seconds):
        """快进/快退"""
        new_time = self.current_time + seconds
        new_time = max(0, min(new_time, self.max_time))
        self.current_time = new_time
        self.update_view()
        self.update_ui_slider()

    def slider_pressed(self):
        self.was_playing = self.is_playing
        if self.is_playing:
            self.toggle_play()

    def slider_released(self):
        if self.was_playing:
            self.toggle_play()

    def slider_moved(self, val):
        # 只有当非播放状态或用户拖动时才更新
        self.current_time = val / 100.0
        self.lbl_time.setText(f"{self.current_time:.2f} s")
        self.update_view()

    def update_ui_slider(self):
        self.slider.blockSignals(True)
        self.slider.setValue(int(self.current_time * 1000))
        self.slider.blockSignals(False)
        self.lbl_time.setText(f"{self.current_time:.6f} s")

    def set_speed(self, val):
        self.playback_speed = val

    def get_node_pos(self, nid, t):
        data = self.pos_data[nid]
        times = data['times']
        positions = data['pos']
        
        if t <= times[0]: return positions[0]
        if t >= times[-1]: return positions[-1]
        
        idx = np.searchsorted(times, t)
        if idx == 0: return positions[0]
        
        t0, t1 = times[idx-1], times[idx]
        p0, p1 = positions[idx-1], positions[idx]
        alpha = (t - t0) / (t1 - t0)
        return p0 + (p1 - p0) * alpha

    def update_frame(self):
        step = (1.0 / self.fps) * self.playback_speed
        self.current_time += step
        if self.current_time > self.max_time:
            self.current_time = 0 
            self.toggle_play() # 播完停止
            
        self.update_view()
        self.update_ui_slider()

    def update_view(self):
        curr_node_pos = {}
        
        # 1. 更新节点位置
        for nid, item in self.node_items.items():
            pos = self.get_node_pos(nid, self.current_time)
            curr_node_pos[nid] = pos
            
            item.resetTransform()
            item.translate(pos[0], pos[1], pos[2])
            item.scale(20.0, 20.0, 20.0)

            # 更新文字位置
            if nid in self.node_labels:
                label_pos = np.array([pos[0], pos[1], pos[2] - 20], dtype=np.float32)
                self.node_labels[nid].setData(pos=label_pos)

        # 2. 更新通信
        active_comms = [p for p in self.comm_data if p['send_time'] <= self.current_time <= p['arrive_time']]
        current_pids = set()

        for comm in active_comms:
            pid = comm['id']
            current_pids.add(pid)
            
            src = curr_node_pos[comm['src']]
            dst = curr_node_pos[comm['dst']]
            
            dur = comm['arrive_time'] - comm['send_time']
            prog = (self.current_time - comm['send_time']) / dur if dur > 0 else 1.0
            pkt_pos = src + (dst - src) * prog
            
            if pid not in self.active_packets:
                # 连线
                line = gl.GLLinePlotItem(width=2, color=(1, 1, 1, 0.4))
                # 数据球 (红色)
                ball = gl.GLScatterPlotItem(size=20, color=(1, 0, 0, 1), pxMode=True)
                # 文本
                # 注意：GLTextItem 在 PyQTGraph 中有时会有显示问题
                # 关键：设置为 position 属性，并保证颜色不为黑
                text = gl.GLTextItem(pos=pkt_pos, text=comm['tag'], color=(255, 255, 0, 255))
                # 设置字体大小 (虽然 GLTextItem 缩放通常通过 OpenGL 矩阵，但这能影响清晰度)
                # font = text.font()
                # font.setPointSize(12) 
                # text.setFont(font)
                
                self.view.addItem(line)
                self.view.addItem(ball)
                self.view.addItem(text)
                self.active_packets[pid] = {'line': line, 'ball': ball, 'text': text}
            
            objs = self.active_packets[pid]
            objs['line'].setData(pos=np.array([src, dst]))
            objs['ball'].setData(pos=np.array([pkt_pos]))
            
            # 更新文本位置 (抬高Z轴防止与球重叠)
            text_pos = pkt_pos + np.array([0, 0, 15]) 
            objs['text'].setData(pos=text_pos)

        # 3. 清理过期包
        for pid in list(self.active_packets.keys()):
            if pid not in current_pids:
                objs = self.active_packets[pid]
                self.view.removeItem(objs['line'])
                self.view.removeItem(objs['ball'])
                self.view.removeItem(objs['text'])
                del self.active_packets[pid]

if __name__ == "__main__":
    app = QApplication(sys.argv)
    if len(sys.argv) > 1:
        fname = sys.argv[1]
    else:
        fname = "simulation_log.json"
        import os
        if not os.path.exists(fname):
            print("请先运行 simulation.py 生成数据")
            sys.exit(0)
            
    win = VisualizerWindow(fname)
    win.show()
    sys.exit(app.exec())